import librosa
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

"""
基于重复模式提取的盲源分离demo
"""


class similar_mat_denoise:
    def __init__(self,y,sr):
        '''
        输入48khz的时序信号
        :param y:
        :param sr:
        '''

        self.y = y
        self.sr = sr
        self.n_fft = int(25 / 1000 * self.sr)         # 窗口长度,业界常见的帧长度是25ms
        self.hop_length = int(10 / 1000 * self.sr)    # 窗口重叠数，业界常见的帧间隔是10ms
    def denoise(self):
        y = self.y
        n_fft = self.n_fft
        hop_length = self.hop_length
        noise_stft = librosa.stft(y,n_fft = n_fft,hop_length= hop_length,window = "hamm")  # stft得到时频谱
        noise_stft_real, noise_stft_phase = librosa.magphase(noise_stft, power=1)  # 计算复数图谱的幅度值和相位值

        '''找出每个频率的中位数(论文法)'''
        noise_stft_trans = np.transpose(noise_stft.copy()).real             # 将noise_stft转置并取实数值
        similar_matrix = cosine_similarity(noise_stft_trans)                # 计算相似矩阵
        num_of_maxs = int(0.8*len(similar_matrix))
        V_sim_complex = []
        import heapq
        for idx1,row in enumerate(similar_matrix):         # 第idx1帧，row是第idx1帧的相似度向量
            row = row.tolist()          # array转list
            idx = list(map(row.index, heapq.nlargest(num_of_maxs,row)))  # 求最大的num_of_maxs个索引,nsmallest与nlargest相反，求最小
            similar_frame = noise_stft[:,idx]                  # 获得与取出帧相似度最大的若干帧
            V_sim = np.median(similar_frame, axis=1)           # np.real(noise_stft)每行的中位数。axis=1对应行，axis=0对应列
            temp = []
            for idx2,V in enumerate(V_sim):
                if abs(noise_stft[idx2,idx1].real) <= abs(V.real):
                    temp.append(noise_stft[idx2,idx1])
                else:
                    temp.append(V)
            V_sim_complex.append(temp)
        V_sim_complex = np.transpose(V_sim_complex)

        # '''找出每个频率的中位数(拔萝卜法)'''
        # V_sim = np.median(noise_stft_real, axis=1)  # np.real(noise_stft)每行的中位数。axis=1对应行，axis=0对应列
        # '''重复谱更新部分'''
        # V_sim_complex = []  # 创建空的重复谱，用于保存复数幅值
        # for idx, V in enumerate(V_sim):
        #     temp = []
        #     for x in np.real(noise_stft[idx]):
        #         # if abs(x) <= V:
        #         #     temp.append(x)
        #         if abs(x-V) <= 2*abs(np.mean(noise_stft[idx])-V):       # 若当前数值与中值的偏差小于等于均值与中值的偏差，则认为不是瞬态干扰，应保留当前数值
        #             temp.append(x)
        #         else:
        #             temp.append(V)
        #     V_sim_complex.append(temp)
        # V_sim_complex = np.array(V_sim_complex)

        # '''找出每个频率的中位数(夹板法)'''
        # V_sim = np.median(noise_stft_real, axis=1)  # np.real(noise_stft)每行的中位数。axis=1对应行，axis=0对应列
        # '''重复谱更新部分'''
        # V_sim_complex = []  # 创建空的重复谱，用于保存复数幅值
        # for idx, V in enumerate(V_sim):
        #     temp = []
        #     for x in np.real(noise_stft[idx]):
        #         # if abs(x) <= V:
        #         #     temp.append(x)
        #         if x <= 1.2*V and x >= 0.8*V:       # 若当前数值与中值的偏差小于等于均值与中值的偏差，则认为不是瞬态干扰，应保留当前数值
        #             temp.append(x)
        #         else:
        #             temp.append(V)
        #     V_sim_complex.append(temp)
        # V_sim_complex = np.array(V_sim_complex)

        # '''找出每个频率的中位数(平滑法)'''
        # V_sim = np.median(noise_stft_real, axis=1)  # np.real(noise_stft)每行的中位数。axis=1对应行，axis=0对应列
        # '''重复谱更新部分'''
        # V_sim_complex = []  # 创建空的重复谱，用于保存复数幅值
        # from scipy.signal import savgol_filter
        # for idx, V in enumerate(V_sim):
        #     temp = savgol_filter(noise_stft[idx], window_length=51, polyorder=2)
        #     V_sim_complex.append(temp)
        # V_sim_complex = np.array(V_sim_complex)

        '''计算维纳滤波系数部分'''
        # H_sim = V_sim_complex / noise_stft  # 维纳滤波系数,H_sim是复数矩阵

        '''哈曼达积时频掩蔽'''
        # back_spec = noise_stft * H_sim  # 背景时频谱（振动信号），复数矩阵
        back_spec = V_sim_complex
        noise_spec = noise_stft - back_spec  # 噪声信号，复数矩阵

        '''ISTFT部分'''
        istft_y = librosa.istft(noise_stft,n_fft = n_fft,hop_length= hop_length,window = "hamm")
        istft_back_spec = librosa.istft(back_spec,n_fft = n_fft,hop_length= hop_length,window = "hamm")
        istft_noise_spec = librosa.istft(noise_spec, n_fft=n_fft, hop_length=hop_length, window="hamm")

        '''画图部分'''
        # fig = plt.figure()
        # ax1 = fig.add_subplot(121)  # 画原始时域信号
        # ax2 = fig.add_subplot(122)  # 画降噪时域信号
        # ax1.plot(y)
        # ax2.plot(istft_back_spec)
        # plt.show()
        return noise_stft, back_spec, noise_spec, istft_y,istft_back_spec,istft_noise_spec      # 含噪时频、背景时频、噪声时频、含噪时域、背景时域、噪声时域




if __name__ == '__main__':
    from data_pre_processing import Universal_tool
    filepath = '../data/wav_data_V4/normal/环境声响_正常运行_normal_01.wav'
    # filepath = './data/wav_data_unknown/大喊2.wav'

    y,sr = librosa.load(filepath)

    # y = Universal_tool.AVF(y,sr).astype('float')    # 幅值方差波动去噪

    zqd = similar_mat_denoise(y,sr)
    data_stft,back_spec,noise_spec,istft_y,istft_back_spec,istft_noise_spec = zqd.denoise()
    mel_spec = librosa.feature.melspectrogram(y=y, sr=zqd.sr, n_fft=zqd.n_fft, hop_length=zqd.hop_length)  # Mel谱
    sr = zqd.sr
    fig = plt.figure()

    '''画Mel谱与时频谱的对比图'''
    plt.subplot(1, 2, 1)
    plt.title("Mel spectrum")
    librosa.display.specshow(np.log(abs(mel_spec)), sr=zqd.sr,
                             n_fft=zqd.n_fft, hop_length=zqd.hop_length, x_axis='time',y_axis='mel')
    plt.colorbar(format='%+2.0f dB')
    plt.subplot(1, 2, 2)
    plt.title("Time-frequency spectrum")
    librosa.display.specshow(np.log(abs(data_stft)), sr=zqd.sr,
                             n_fft=zqd.n_fft, hop_length=zqd.hop_length, x_axis='time', y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    fig.tight_layout(pad=0.1, w_pad=0, h_pad=0.4)
    plt.show()

    '''画盲源分离对比图'''
    plt.subplot(3,1,1)
    plt.title("mixed_spec")
    librosa.display.specshow(np.log(abs(data_stft)), sr = sr,
                             n_fft=zqd.n_fft, hop_length=zqd.hop_length,x_axis='time',y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    plt.subplot(3,1,2)
    plt.title("back_spec")
    librosa.display.specshow(np.log(abs(back_spec)), sr = sr,
                             n_fft=zqd.n_fft, hop_length=zqd.hop_length,x_axis='time',y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    plt.subplot(3,1,3)
    plt.title("noise_spec")
    librosa.display.specshow(np.log(abs(noise_spec)), sr = sr,
                             n_fft=zqd.n_fft, hop_length=zqd.hop_length,x_axis='time',y_axis='hz')
    plt.colorbar(format='%+2.0f dB')
    fig.tight_layout(pad=0.1, w_pad=0, h_pad=0.4)
    plt.show(block = True)

    '''生成盲源分离后的wav文件'''
    Universal_tool.array2wav(istft_y,'istft_y.wav',sr)
    Universal_tool.array2wav(istft_back_spec, 'istft_back_spec.wav', sr)
    Universal_tool.array2wav(istft_noise_spec, 'istft_noise_spec.wav', sr)

















